package com.example.sharding.jdbc.demo.config;

import com.example.sharding.jdbc.demo.util.MonthPreciseShardingUtil;
import com.example.sharding.jdbc.demo.constants.ShardingConstants;
import com.example.sharding.jdbc.demo.util.DateUtil;
import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.standard.PreciseShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.RangeShardingValue;
import org.apache.shardingsphere.sharding.api.sharding.standard.StandardShardingAlgorithm;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Timestamp;
import java.time.LocalDateTime;
import java.util.Collection;
import java.util.Date;
import java.util.Properties;
import java.util.Set;

/**
 * 按月分片算法
 *
 * @author airtrioa
 * @date 2021/4/30
 */
public class MonthShardingAlgorithm implements StandardShardingAlgorithm<Date> {
  private static final Logger logger = LoggerFactory.getLogger(MonthShardingAlgorithm.class);
  private Properties props = new Properties();


  @Override
  public String doSharding(Collection<String> collection, PreciseShardingValue<Date> preciseShardingValue) {
    // 获取到月份
    Date value = preciseShardingValue.getValue();
    int month = DateUtil.toLocalDateTime(value).getMonthValue();
    // 筛选该次查询属于哪个表
    for (String tableName : collection) {
      // 获取表名的序列号 ->  即表的月份
      String tableIndex = MonthPreciseShardingUtil.obtainTableIndex(tableName);
      // 序列号 包含 01、1两种
      if (Integer.parseInt(tableIndex) == month) {
        logger.info("定位到分片表:[{}]", tableName);
        return tableName;
      }
    }
    throw new IllegalArgumentException("未找到匹配的数据表");
  }

  @Override
  public Collection<String> doSharding(Collection<String> collection, RangeShardingValue<Date> rangeShardingValue) {
    Range<Date> valueRange = rangeShardingValue.getValueRange();
    LocalDateTime lowerEndpoint = DateUtil.toLocalDateTime(valueRange.lowerEndpoint());
    LocalDateTime upperEndpoint = DateUtil.toLocalDateTime(valueRange.upperEndpoint());
    if (valueRange.lowerEndpoint() instanceof Timestamp) {
      lowerEndpoint = ((Timestamp) valueRange.lowerEndpoint()).toLocalDateTime();
      upperEndpoint = ((Timestamp) valueRange.upperEndpoint()).toLocalDateTime();
    }
    Set<String> tables = MonthPreciseShardingUtil.rangeShardingTables(collection, lowerEndpoint, upperEndpoint);
    logger.info("定位到分片表:{}", tables);
    return tables;
  }

  @Override
  public void init() {
    logger.info("初始化按月分片算法");
  }

  @Override
  public String getType() {
    return ShardingConstants.MONTH_SHARD;
  }

  @Override
  public Properties getProps() {
    return props;
  }

  @Override
  public void setProps(Properties props) {
    this.props = props;
  }
}
